###model primitives
@with_kw struct Primitives
    r::Float64 = 0.03 #interest rate
    r_d::Float64 = 0.07 #student loan interest rate
    η::Float64 = 2.0 #CRRA coefficient
    β::Float64 = 1 / (1 + 0.03)
    dbar = 0.735
    cfloor::Float64 = 0.111 #consumption floor
    prob_forb::Float64 = 0.36 #probability of forebearance if unemployed or NILF
    p_min::Float64 = 0.015 #minimum annual debt payment
    p_max::Float64 = 0.105
    thresh_sub::Float64 = 1.4 #parent income threshold for subsidized loan eligibility

    #various borrowing limits
    dbar_nolim::Float64 = 3.5
    dbar_baseline::Float64 = 0.735
    dbar_rec::Float64 = 0.991

    #debt upper limit to use in policy functions
    dcap::Float64 = 1.0
    
    #baseline student loan grid
    d_grid::Vector{Float64} = collect((range(0.0, length = 10, dcap)))
    nd::Int64 = length(d_grid)

    #ability grid
    θ_grid::Vector{Float64} = collect(0:0.25:1.0) #ability grid
    nθ::Int64 = length(θ_grid)

    ###asset grid
    a_grid::Vector{Float64} = collect(range(-0.111, length = 20, 6.25)) #-140000 to 250000, 25 states
    na::Int64 = length(a_grid)
    
    ##age grid
    t_grid::Vector{Int64} = collect(18:1:65) #start at 19, end at 66
    nt::Int64 = length(t_grid)

    ##education grid
    e_grid::Vector{Int64} = collect(1:1:3) #hs, somecoll, coll
    ne::Int64 = length(e_grid)

    ##experience grid
    x_grid::Vector{Float64} = collect(0:1:47) #start to finish
    nx::Int64 = length(x_grid)

    #parent income/wealth grid
    I_grid::Vector{Float64} = collect(range(0.01, length = 30, stop = 5.0))
    nI::Int64 = length(I_grid)

    W_grid::Vector{Float64} = collect(range(0.01, length = 20, stop = 12.5))
    nW::Int64 = length(W_grid)
end

###parameters to estimate
@with_kw struct Estimands
    guess::Vector{Float64} = zeros(30)

    #parent preferences
    ω::Float64 = guess[1] #wealth weight
    α::Float64 = guess[2] #altruism

    #college preferences
    ψ_0::Float64 = guess[3] #fixed psychic cost
    ψ_1::Float64 = guess[4] #ability effect
    ψ_2::Float64 = guess[5] #ability-hq effect
    σ_ζ::Float64 = guess[6] #spread of college preference shocks
    #ψ_3::Float64 = guess[32] #debt intercept
    ψ_3::Float64 = 0.0
    ψ_4::Float64 = guess[32] #debt slope
    #ψ_4::Float64 = 0.0

    #earnings determinants
    γ_0_hs::Float64 = guess[7] #constant, hs
    γ_1_hs::Float64 = guess[8] #ability, hs
    γ_2_hs::Float64 = guess[9] #exp, hs
    γ_3_hs::Float64 = guess[10] #exp2, hs
    γ_0_sc::Float64 = guess[11] #constant, sc
    γ_1_sc::Float64 = guess[12] #ability, sc
    γ_2_sc::Float64 = guess[13] #exp, sc
    γ_3_sc::Float64 = guess[14] #exp2, sc
    γ_0_coll::Float64 = guess[15] #constant, coll
    γ_1_coll::Float64 = guess[16] #ability, coll
    γ_2_coll::Float64 = guess[17] #exp, coll
    γ_3_coll::Float64 = guess[18] #exp2, coll
    γ_grid::Array{Float64} = [γ_0_hs γ_1_hs γ_2_hs γ_3_hs; γ_0_sc γ_1_sc γ_2_sc γ_3_sc; γ_0_coll γ_1_coll γ_2_coll γ_3_coll] #all togother!

    #leisure preferences
    λ_0::Float64 = guess[19] #constant
    λ_1::Float64 = guess[20] #ability
    λ_2::Float64 = guess[21] #age
    λ_3::Float64 = guess[22] #age2
    #λ_4::Float64 = guess[32] #age3
    λ_4::Float64 = 0.0 #age3
    ξ::Float64 = guess[23] #switch cost
    σ_l::Float64 = guess[24] #shock spread
    
    #employment determinants
    π_0::Float64 = guess[25] #constant
    π_1::Float64 = guess[26] #previously employed
    π_2::Float64 = guess[27] #some college
    π_3::Float64 = guess[28] #college
    π_4::Float64 = guess[29] #exp
    π_5::Float64 = guess[30] #exp*sc
    π_6::Float64 = guess[31] #exp*coll
end

###model results
mutable struct Results
    val_func_emp::SharedArray{Float64, 6} #value function, employed
    pol_func_emp::SharedArray{Float64, 6} #policy function, employed
    
    val_func_unemp::SharedArray{Float64, 7} #value function, unemployed
    pol_func_unemp::SharedArray{Float64, 7} #policy function, unemployed

    val_func_nilf::SharedArray{Float64, 7} #value function, unemployed
    pol_func_nilf::SharedArray{Float64, 7} #policy function, unemployed

    pol_func_p_I::SharedArray{Float64, 3} #parent policy function, income
    pol_func_p_W::SharedArray{Float64, 3} #parent policy function, wealth
    pol_func_p_τ::SharedArray{Float64, 3} #parent policy function, total
end

###initialize model primitives and value/policy functions to fill
function Initialize(guess::Vector{Float64}, utilities::Vector{Matrix{Float64}}; dlim::Int64 = 0, decomp::Int64 = 0, rec_pol::Int64 = 0)
    
    d_opts = utilities[9]
    dcap = 1.0 #baseline debt upper bound to use in state space
    p_max = 0.15 #baseline maximum payment

    #update to recession borring limit if that's where we are
    if rec_pol == 1 && (decomp == 0 || decomp == 3)
        dcap = 1.3
        p_max = 0.2
    end

    #switch if we've hard-coded the borrowing limit
    if dlim!=0
        dcap = d_opts[dlim, 2]
        p_max = d_opts[dlim, 3]
    end
        
    prim = Primitives(dcap = dcap, p_max = p_max) #boot up primitives
    est = Estimands(guess = guess) #pass guess of estimands

    ###initialize results struct##
    val_func_emp = SharedArray{Float64}(prim.nθ, prim.ne, prim.nx, prim.na, prim.nt, prim.nd)
    pol_func_emp = SharedArray{Float64}(prim.nθ, prim.ne, prim.nx, prim.na, prim.nt, prim.nd)
    val_func_unemp = SharedArray{Float64}(prim.nθ, prim.ne, prim.nx, prim.na, prim.nt, prim.nd, 2)
    pol_func_unemp = SharedArray{Float64}(prim.nθ, prim.ne, prim.nx, prim.na, prim.nt, prim.nd, 2)
    val_func_nilf = SharedArray{Float64}(prim.nθ, prim.ne, prim.nx, prim.na, prim.nt, prim.nd, 2)
    pol_func_nilf = SharedArray{Float64}(prim.nθ, prim.ne, prim.nx, prim.na, prim.nt, prim.nd, 2)

    #parent stuff
    pol_func_p_I = SharedArray{Float64}(prim.nI, prim.nW, prim.nθ)
    pol_func_p_W = SharedArray{Float64}(prim.nI, prim.nW, prim.nθ)
    pol_func_p_τ = SharedArray{Float64}(prim.nI, prim.nW, prim.nθ)

    res = Results(val_func_emp, pol_func_emp, val_func_unemp, pol_func_unemp,
    val_func_nilf, pol_func_nilf, pol_func_p_I, pol_func_p_W, pol_func_p_τ)

    prim, est, res #return deliverables
end

###function that solves the model
function Solve_model(guess::Vector{Float64}, utilities::Vector{Matrix{Float64}}; 
    rec::Int64 = 0, cfact::Int64 = 0, pbr::Int64 = 0, dlim::Int64 = 0, decomp::Int64 = 0, rec_pol::Int64 = 0)
    
    prim, est, res = Initialize(guess, utilities; dlim = dlim, decomp = decomp, rec_pol = rec_pol)
    @unpack t_grid, nt, nθ, ne = prim

    @sync @distributed for i_θ = 1:nθ
        for i_e = 1:ne
            for i = 1:nt #loop over ages
                i_t = nt - i + 1 #now backward -- storing index of age state
                Update_valfunc(prim, est, res, i_t, i_θ, i_e, utilities; rec=rec, decomp=decomp) #call function that updates value functions!
            end
        end
    end

    Valfunc_parent(prim, est, res, utilities; rec=rec, cfact=cfact, pbr=pbr, dlim=dlim, decomp=decomp, rec_pol=rec_pol)
    
    ##return deliverables
    prim, est, res
end

###equal-mass discretization of continuous dsitribution method from Kennan (2006)
function discretize_lognormal(μ::Float64, σ::Float64, n::Int64)
    dist = LogNormal(μ, σ) #set up log normal distribution
    shocks = zeros(n) #preallocate market luck states

    for i = 1:n #begin loop to fill discretized vector
        quant = (2*i-1)/(2*n) #desired quantile
        shocks[i] = quantile(dist, quant) #fill element of state vector
    end
    shocks #return deliverable
end

###equal-mass discretization of continuous dsitribution method from Kennan (2006)
function discretize_normal(μ::Float64, σ::Float64, n::Int64)
    dist = Normal(μ, σ) #set up log normal distribution
    shocks = zeros(n) #preallocate market luck states

    for i = 1:n #begin loop to fill discretized vector
        quant = (2*i-1)/(2*n) #desired quantile
        shocks[i] = quantile(dist, quant) #fill element of state vector
    end
    shocks #return deliverable
end


###utility (CRRA)
function utility(c::Float64, η::Float64)
    util = (c^(1-η) - 1) / (1-η)
    if c<=0
        util = -Inf #none of this nonsense!
    end
    util #return
end

###debt payments
function payment(d::Float64, r::Float64, t::Int64)
    payment = d * (r) / (1-(1+r)^(-t))
    payment
end

##grant calculator
function Grants(I::Float64, s::Int64, sched::Array{Float64,2}, W::Float64; cfact::Int64 = 0)
    grant, grant_fed = 0, 0

    if s!=0 #do stuff
        
        #row of federal and state aid schedules
        row_fed = s
        row_state = s + 4

        ##convert income to 10k 2012 to get the g5 flag
        inc_2000 = I
        inc_2000 *=4
        inc_2000 *=0.78235 #now in 10000 2012

        grant_fed = sched[row_fed, 2] +  sched[row_fed, 3]*(I*4) + sched[row_fed, 4]*(inc_2000>5) + sched[row_fed, 5]*(inc_2000>5) * (I*4)
        grant_state = sched[row_state, 2] +  sched[row_state, 3]*(I*4) + sched[row_state, 4]*(inc_2000>5) + sched[row_state, 5]*(inc_2000>5) * (I*4)
        grant = grant_fed + grant_state
    end

    #update federal grant if we're in a particular policy counterfactual 
    if cfact == 1 && I<=3.125 #Pell expansion: everybody under 125k just gets 15k 
        grant -= grant_fed
        grant += 0.375 #flat 15k
        grant_fed = 0.375 #flat 15k
    end

    #expand Pell generosity by a factor of 2
    if cfact == 2 
        grant += grant_fed
        grant_fed *= 2
    end

    #final cfact: wealth-dependent Pell expansion
    if cfact == 3 
        grant += (0.5 - 0.2*W + 0.2*W*(W>2.5) - 0.5*(W>2.5))
        grant_fed += (0.5 - 0.2*W + 0.2*W*(W>2.5) - 0.5*(W>2.5))
    end

    grant, grant_fed
end

###merit-based grants
function Grants_merit(θ::Float64, s::Int64, sched::Array{Float64, 2})
    #initialize
    grant_merit = 0

    #do stuff
    if s!=0
        row_select = s
        grant_merit = sched[row_select, 2] + θ * sched[row_select, 3] #linear function of ability
    end

    #return
    grant_merit
end

###function for computing graduation probabilties based on college selection and level of ability
function Grad_prob(θ::Float64, s::Int64, probs::Array{Float64, 2})
    prob_drop, prob_sc, prob_coll = 0, 0, 0 #initialize -- probabilities of 

    #update if actually went somewhere
    row_drop = (s-1)*3 + 1
    row_sc = (s-1)*3 + 2

    prob_drop = probs[row_drop, 3] + θ * probs[row_drop, 4]
    prob_sc = probs[row_sc, 3] + θ * probs[row_sc, 4]
    prob_coll = 1 - prob_drop - prob_sc
    prob_drop, prob_sc, prob_coll
end

#function for getting interpolated index of a given Float with an arbitrary grid, assuming linear interpolation
function get_index(val::Float64, grid::Array{Float64,1})
    n = length(grid)
    index = 0 #preallocation
    if val<=grid[1] #LEQ smallest element
        index = 1
    elseif val>=grid[n] #GEQ biggest element
        index = n
    else
        index_upper = findfirst(z->z>val, grid)
        index_lower = index_upper - 1
        val_upper, val_lower = grid[index_upper], grid[index_lower] #values
        index = index_lower + (val - val_lower)  / (val_upper - val_lower) #weighted average
    end
    index #return
end

####net taxes given income
function tax(i::Float64)
    #note: λ here was kinda just chosen to get the right average
    hsv_τ, hsv_λ = 0.1807746, 6.2659073
    transfer_all = 2300
    i*=40000

    tax = i - hsv_λ * i^(1-hsv_τ) #hsv tax function
    tax -= transfer_all #universal transfer/deduction lump
    tax/40000 #return
end

#standard normal CDF
function 𝚽(x::Float64)
    return cdf(Normal(), x)
end

#mills ratio
function Mills(x::Float64, μ::Float64, σ::Float64)
    frac = (x - μ)/σ
    mills = μ + σ * pdf(Normal(), frac)/(1-cdf(Normal(), frac))
    mills #return
end
############
